{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Serving PyTorch Models In Production With Amazon SageMaker and TorchServe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Your Hosting Environment\n",
    "The focus of this lab is around model serving. In that vain, we have taken care of of the data preparation and model training. \n",
    "This lab exercise is using a [HuggingFace Transformer](https://huggingface.co/transformers/) which provides us with a general-purpose architecture for Natural Language Understanding (NLU). Specifically, we are presenting you with a [RoBERTa base](https://huggingface.co/roberta-base) transformer that was fined tuned to perform sentiment analysis. The pre-trained checkpoint loads the additional head layers and will output ``positive``, ``neutral``, and ``negative`` sentiment or text. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "import pandas as pd\n",
    "\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.utils import name_from_base\n",
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "\n",
    "from sagemaker.predictor import Predictor\n",
    "from sagemaker.serializers import JSONLinesSerializer\n",
    "from sagemaker.deserializers import JSONLinesDeserializer\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "bucket = sess.default_bucket()\n",
    "role = sagemaker.get_execution_role()\n",
    "region = boto3.Session().region_name\n",
    "\n",
    "sm = boto3.Session().client(service_name=\"sagemaker\", region_name=region)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r training_job_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(training_job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store -r transformer_pytorch_model_dir_s3_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(transformer_pytorch_model_dir_s3_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Your Endpoint\n",
    "We will now create and deploy our model. To begin, we need to construct a new PyTorchModel object which points to the pre-trained model artifacts from the above step and also points to the inference code that we wish to use. We will then call the deploy method to launch the deployment container on our TorchServe powered Amazon SageMaker endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StarRatingPredictor(Predictor):\n",
    "    def __init__(self, endpoint_name, sagemaker_session):\n",
    "        super().__init__(\n",
    "            endpoint_name,\n",
    "            sagemaker_session=sagemaker_session,\n",
    "            serializer=JSONLinesSerializer(),\n",
    "            deserializer=JSONLinesDeserializer(),\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "timestamp = int(time.time())\n",
    "\n",
    "pytorch_model_name = \"{}-{}-{}\".format(training_job_name, \"pt\", timestamp)\n",
    "\n",
    "print(pytorch_model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = PyTorchModel(\n",
    "    model_data=transformer_pytorch_model_dir_s3_uri + \"model.tar.gz\",\n",
    "    name=pytorch_model_name,\n",
    "    role=role,\n",
    "    entry_point=\"inference.py\",\n",
    "    source_dir=\"code-pytorch\",\n",
    "    framework_version=\"1.6.0\",\n",
    "    py_version=\"py3\",\n",
    "    predictor_cls=StarRatingPredictor,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "pytorch_endpoint_name = \"{}-{}-{}\".format(training_job_name, \"pt\", timestamp)\n",
    "\n",
    "print(pytorch_endpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = model.deploy(\n",
    "    initial_instance_count=1, instance_type=\"ml.m4.xlarge\", endpoint_name=pytorch_endpoint_name, wait=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(predictor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "\n",
    "display(\n",
    "    HTML(\n",
    "        '<b>Review <a target=\"blank\" href=\"https://console.aws.amazon.com/sagemaker/home?region={}#/endpoints/{}\">SageMaker REST Endpoint</a></b>'.format(\n",
    "            region, pytorch_endpoint_name\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "waiter = sm.get_waiter(\"endpoint_in_service\")\n",
    "waiter.wait(EndpointName=pytorch_endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# _Wait Until the ^^ Endpoint ^^ is Deployed_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pytorch_endpoint_arn = sm.describe_endpoint(EndpointName=pytorch_endpoint_name)[\"EndpointArn\"]\n",
    "print(pytorch_endpoint_arn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.lineage.visualizer import LineageTableVisualizer\n",
    "\n",
    "lineage_table_viz = LineageTableVisualizer(sess)\n",
    "lineage_table_viz_df = lineage_table_viz.show(endpoint_arn=pytorch_endpoint_arn)\n",
    "lineage_table_viz_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Perform Predictions With A TorchServe Backend Amazon SageMaker Endpoint\n",
    "Here, we will pass sample strings of text to the endpoint in order to see the sentiment. We give you one example of each, however, feel free to play around and change the strings yourself! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "inputs = [{\"features\": [\"This is great!\"]}, {\"features\": [\"This is bad.\"]}]\n",
    "\n",
    "predicted_classes = predictor.predict(inputs)\n",
    "\n",
    "for predicted_class in predicted_classes:\n",
    "    print(\"Predicted star_rating: {}\".format(predicted_class))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predict the `star_rating` with `review_body` Samples from our TSV's\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import pandas as pd\n",
    "\n",
    "df_reviews = pd.read_csv(\n",
    "    \"./data/amazon_reviews_us_Digital_Software_v1_00.tsv.gz\",\n",
    "    delimiter=\"\\t\",\n",
    "    quoting=csv.QUOTE_NONE,\n",
    "    compression=\"gzip\",\n",
    ")\n",
    "\n",
    "df_sample_reviews = df_reviews[[\"review_body\", \"star_rating\"]].sample(n=50)\n",
    "df_sample_reviews = df_sample_reviews.reset_index()\n",
    "df_sample_reviews.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "\n",
    "def predict(review_body):\n",
    "    inputs = [{\"features\": [review_body]}]\n",
    "    predicted_classes = predictor.predict(inputs)\n",
    "    return predicted_classes[0][\"predicted_label\"]\n",
    "\n",
    "\n",
    "df_sample_reviews[\"predicted_class\"] = df_sample_reviews[\"review_body\"].map(predict)\n",
    "df_sample_reviews.head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pass Variables to the Next Notebook(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store pytorch_endpoint_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Release Resources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sm.delete_endpoint(\n",
    "#     EndpointName=pytorch_endpoint_name\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%html\n",
    "\n",
    "<p><b>Shutting down your kernel for this notebook to release resources.</b></p>\n",
    "<button class=\"sm-command-button\" data-commandlinker-command=\"kernelmenu:shutdown\" style=\"display:none;\">Shutdown Kernel</button>\n",
    "        \n",
    "<script>\n",
    "try {\n",
    "    els = document.getElementsByClassName(\"sm-command-button\");\n",
    "    els[0].click();\n",
    "}\n",
    "catch(err) {\n",
    "    // NoOp\n",
    "}    \n",
    "</script>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript\n",
    "\n",
    "try {\n",
    "    Jupyter.notebook.save_checkpoint();\n",
    "    Jupyter.notebook.session.delete();\n",
    "}\n",
    "catch(err) {\n",
    "    // NoOp\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Internal - DO NOT RUN - WILL REMOVE SOON"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%bash\n",
    "\n",
    "# aws sagemaker-runtime invoke-endpoint \\\n",
    "#     --endpoint-name \"tensorflow-training-2021-01-28-01-19-50-987-pt-1611813221\" \\\n",
    "#     --content-type application/jsonlines \\\n",
    "#     --accept application/jsonlines \\\n",
    "#     --body $'{\"features\":[\"Amazon gift cards are the best\"]}\\n{\"features\":[\"It is the worst\"]}' >(cat) 1>/dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !rm model.tar.gz\n",
    "# !aws s3 cp s3://sagemaker-us-east-1-835319576252/tensorflow-training-2021-01-28-01-19-50-987/output/model.tar.gz ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !rm -rf ./model\n",
    "# !mkdir -p  ./model\n",
    "# !tar -xvzf ./model.tar.gz -C model/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !cp ./code/inference.py model/code/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !cat model/code/inference.py"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
